import cv2
import time
import os
from ultralytics import YOLOv10

def Yolov10Detector(frame, model, image_size, conf_threshold):
    results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold)
    frame = results[0].plot()
    return frame, len(results[0].boxes) > 0  # 返回处理后的帧和是否有检测结果

def main():
    # 参数设置
    image_size = 640
    conf_threshold = 0.1
    model = YOLOv10("models/YOLOv10-FireSmoke-M.pt")
    video_path = "1.mp4"  # 改为您的视频文件路径
    
    # 输出图片设置
    output_dir = "detected_frames"
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"打开视频文件失败: {video_path}")
        return
    
    # 获取视频参数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    print(f"开始处理视频: {video_path}")
    print(f"总帧数: {total_frames}")
    print("按Q键可提前终止处理...")

    saved_count = 0
    last_save_time = 0  # 上次保存时间
    save_interval = 5   # 5秒保存一次
    
    while True:
        success, frame = cap.read()
        start_time = time.time()

        if not success:
            print("\n视频处理完成！")
            break
            
        current_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
        current_time = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000  # 当前时间(秒)
        
        # 执行检测
        processed_frame, has_detection = Yolov10Detector(frame, model, image_size, conf_threshold)
        
        # 计算并显示FPS
        end_time = time.time()
        current_fps = 1 / (end_time - start_time)
        framefps = "FPS:{:.2f}".format(current_fps)
        cv2.rectangle(processed_frame, (10, 1), (120, 20), (0, 0, 0), -1)
        cv2.putText(processed_frame, framefps, (15, 17), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
        
        # 保存检测到的帧（每5秒保存一次）
        if has_detection and (current_time - last_save_time) >= save_interval:
            output_path = os.path.join(output_dir, f"frame_{current_frame:06d}_{current_time:.1f}s.jpg")
            cv2.imwrite(output_path, processed_frame)
            saved_count += 1
            last_save_time = current_time
            print(f"保存检测结果: {output_path}")
        
        # 显示处理进度
        print(f"处理进度: {current_frame}/{total_frames} 帧 (已保存: {saved_count})", end="\r")
        
        cv2.imshow("yolov10-视频识别", processed_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            print("\n用户提前终止处理")
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    print(f"\n处理完成！共保存 {saved_count} 张检测结果到目录: {output_dir}")

if __name__ == "__main__":
    main()